#!/usr/bin/perl
#
# getRandomLines.pl - brendan veeneman 5/2012
#
# N - choose N chunks randomly
# K - from K total chunks
# C - chunk size.  Defaults to 1 (one line). 2 = fa, 4 = fq
#
use List::Util 'shuffle';
use Getopt::Long;
use strict;

#global variables
my ($N,$K,$infile1,$infile2,$outfile1,$outfile2,$handle1,$handle2);
my $C = 1;

sub parseArgs()
{  
  GetOptions(
    'N=i' => \$N,
    'K:i' => \$K,
    'C:i' => \$C,
    '1=s' => \$infile1,
    '2:s' => \$infile2,
    'out1:s' => \$outfile1,
    'out2:s' => \$outfile2);
  
  die "getRandomLines.pl
  Description:
    get N random \"chunks\" of arbitrary size from one
    or two (paired) input files, without replacement.

  Arguments:
    [-1 infp1]     : input file 1
    (-2 infp2)     : input file 2
    (-out1 outfp1) : output file 1 (default: stdout)
    (-out2 outfp2) : output file 2 (default: stdout)
    (-C C)         : chunk size    (default: 1 line)
    [-N N]         : choose N chunks
    (-K K)         : total number of chunks\n"
    if(!$N || !$infile1 || $C < 1);
  
  die "$infile1 unreadable." if(!(-r $infile1));
  die "$infile2 unreadable." if($infile2 && !(-r $infile2));
  die "$outfile1 unwriteable." if($outfile1 && !(-w $outfile1));
  die "$outfile2 unwriteable." if($outfile2 && !(-w $outfile2));
  
  $K = (scalar(`wc -l $infile1`) / $C) if(!$K); #chunks
  
  die "N must be <= K" if($N > $K);
}

parseArgs();

#generate the 1 to K sequence, and randomize it.  then, take
# the first N elements, and finally, sort those
# ps - they're 1 indexed
my @seqs = sort {$a <=> $b} ((shuffle(1..$K))[0..($N - 1)]);

my $seq_num = shift @seqs;
open(IFP1,$infile1);

if($outfile1) { open($handle1,$outfile1); }
else { $handle1 = *STDOUT; }
if($outfile2) { open($handle2,$outfile2); }
else { $handle2 = *STDOUT; }

if(!$infile2)
{
  for(my $i = 1; $i <= $K; $i++)
  {
    if($i == $seq_num) #we want this chunk
    {
      for(my $j = 0; $j < $C; $j++)
      {
        print $handle1 scalar <IFP1>;
      }
      $seq_num = shift @seqs;
      last if(!defined($seq_num));
    }
    else #we dont want this chunk
    {
      for(my $j = 0; $j < $C; $j++)
      {
        <IFP1>;
      }
    }
  }
}
else #replicated code saves many if statements
{
  open(IFP2,$infile2);
  for(my $i = 1; $i <= $K; $i++)
  {
    if($i == $seq_num) #we want this chunk
    {
      for(my $j = 0; $j < $C; $j++)
      {
        print $handle1 scalar <IFP1>;
      }
      for(my $j = 0; $j < $C; $j++)
      {
        print $handle2 scalar <IFP2>;
      }
      
      $seq_num = shift @seqs;
      last if(!defined($seq_num));
    }
    else #we dont want this chunk
    {
      for(my $j = 0; $j < $C; $j++)
      {
        <IFP1>; <IFP2>;
      }
    }
  }
}